import torch.nn as nn
from torch.nn.utils import rnn

import seq2seq.config as config


class Encoder(nn.Module):
    def __init__(self):
        super(Encoder, self).__init__()
        self.embedding = nn.Embedding(num_embeddings=len(config.num_sequence), embedding_dim=config.embedding_dim,
                                      padding_idx=config.num_sequence.PAD)

        self.gru = nn.GRU(input_size=config.embedding_dim, batch_first=True, num_layers=config.num_layers,
                          hidden_size=config.hidden_size)

    def forward(self, input, input_length):
        embedded = self.embedding(input)
        embedded = rnn.pack_padded_sequence(embedded, lengths=input_length, batch_first=True)  # 打包
        output, hidden = self.gru(embedded)

        out, out_length = rnn.pad_packed_sequence(output, batch_first=True, padding_value=config.num_sequence.PAD)  # 解包
        return out, hidden

